package annTrain;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Random;

import javax.imageio.ImageIO;

import aNN.InputTargetPair;

import projectUtils.ByteUtil;
import projectUtils.ImageArrayUtil;
import projectUtils.ImageDistortionUtil;
import projectUtils.ImageLoader;

public class GenericImageTrainingSet implements TrainingSet {

	private static Random randomGen = new Random();
	HashMap<String, List<String>> imageFilePaths;
	List<Integer> sequence;
	List<String> sequenceOutput;
	Integer index;
	private int numOutput;
	private int tokenWidth, tokenHeight;
	private boolean includeNegativeEntries;
	private boolean nullToggler = true;
	
	public GenericImageTrainingSet(int numOutput, int tokenWidth, int tokenHeight, boolean includeNegativeEntries)
	{
		this.numOutput = numOutput;
		this.tokenWidth = tokenWidth;
		this.tokenHeight = tokenHeight;
		this.includeNegativeEntries = includeNegativeEntries;
	}
	
	@Override
	public void GenerateRandomSequence(int numRepetition) {
		// TODO Auto-generated method stub
		index = 0;
		sequence = new ArrayList<Integer>(numRepetition * (numOutput * 2));
		sequenceOutput = new ArrayList<String>(numRepetition * (numOutput * 2));
		List<String> totalOutputs = new ArrayList<String>(numRepetition * (numOutput * 2));
		HashMap<String, List<Integer>> tempList = new HashMap<String, List<Integer>>();
		for (String folderName : imageFilePaths.keySet())
		{
			tempList.put(folderName, new ArrayList<Integer>());
			if (folderName.matches("((-|\\+)?[0-9]+(\\.[0-9]+)?)+"))
			{
				for (int j = 0; j < numRepetition; j++)
				{
					totalOutputs.add(folderName);
				}
			}
			else
			{
				//NOTICE THIS AND MAYBE CHANGE IT
				for (int j = 0; j < numRepetition * 0; j++)
				{
					totalOutputs.add(folderName);
				}
			}
			List<String> filePaths = imageFilePaths.get(folderName);
			for (int j = 0; j < filePaths.size(); j++)
				tempList.get(folderName).add(j);
		}
		
		while (totalOutputs.size() > 0)
		{
			int removeIndex = randomGen.nextInt(totalOutputs.size());
			String output = totalOutputs.remove(removeIndex);
			sequenceOutput.add(output);
			if (tempList.get(output).isEmpty())
			{
				List<String> filePaths = imageFilePaths.get(output);
				for (int j = 0; j < filePaths.size(); j++)
					tempList.get(output).add(j);
			}
			sequence.add(tempList.get(output).remove(randomGen.nextInt(tempList.get(output).size())));
		}
	}

	@Override
	public InputTargetPair Next() {
		TrainerHelper helper = new TrainerHelper();
		InputTargetPair inputTargetPair = null;
		if (nullToggler == true)
		{
			double[] target = new double[0];
			BufferedImage image;
			String filePath = imageFilePaths.get(sequenceOutput.get(index)).get(sequence.get(index));
			Integer outputIndex = -1;
			try
			{
				String classificationFolder = filePath.split("\\\\")[filePath.split("\\\\").length - 2];
				String classificationFile = filePath.split("\\\\")[filePath.split("\\\\").length - 1];
				if (classificationFolder.matches("((-|\\+)?[0-9]+(\\.[0-9]+)?)+"))
				//if (classificationFolder.startsWith("NULL") == false)
					outputIndex = Integer.parseInt(classificationFolder);
				else
					outputIndex = -1;
				target = helper.GetOutputTarget(numOutput, outputIndex);
				image = new ImageLoader().LoadImage(filePath);
				if (image == null)
					throw new Exception();
			}
			catch (Exception ex)
			{
				index++;
				return null;
			}
			 
			byte[][] imageArray;
			if (outputIndex != -1)
			{
				image = ImageDistortionUtil.DistortImage(image);
				image = ImageDistortionUtil.Undistort(image);
				imageArray = ImageArrayUtil.NormalizeArrayValue(ImageArrayUtil.ImageToArray(ImageArrayUtil.CreateResizedCopy(image, tokenWidth, tokenHeight)));
			}
			else
			{
				imageArray = ImageArrayUtil.ImageToArray(ImageArrayUtil.CreateResizedCopy(image, tokenWidth, tokenHeight));
			}
			
//			File outputFile = new File(System.getProperty("user.dir") + "\\temp\\" + new Random().nextLong() + ".png");
//			//ImageIO.write(originalImage, "png", outputFile);
//			try {
//				ImageIO.write(ImageArrayUtil.ArrayToImage(imageArray), "png", outputFile);
//			} catch (IOException e) {
//				// TODO Auto-generated catch block
//				e.printStackTrace();
//			}
			
			double[] input = new double[imageArray.length * imageArray[0].length];
			for (int i = 0; i < imageArray.length; i++)
			{
				for (int j = 0; j < imageArray[i].length; j++)
				{
					input[i * imageArray[i].length + j] = /*Math.round*/((double)ByteUtil.ByteToInt(imageArray[i][j]) / 255.0 * 4.0) / 4.0 * 2.0 - 1.0;
				}
			}
			inputTargetPair = new InputTargetPair(input, target);
			image = null;
			imageArray = null;
			index++;
			nullToggler = true;
		}
		else
		{
			double[] input = new double[tokenHeight * tokenWidth];
			for (int i = 0; i < input.length; i++)
			{
				if (randomGen.nextDouble() < 0.5)
				{
					input[i] = randomGen.nextDouble() * 16;
				}
			}
			double[] target = helper.GetOutputTarget(numOutput, -1);
			inputTargetPair = new InputTargetPair(input, target);
			nullToggler = true;
		}
		
		return inputTargetPair;
	}

	@Override
	public boolean HasNext() {
		return index < sequence.size() - 1;
	}

	@Override
	public void LoadTrainingSet(String relativePath, String labelFile,
			String imageFile) {
		imageFilePaths = GetFileList(relativePath, new HashMap<String, List<String>>());
	}
	
	
	private HashMap<String, List<String>> GetFileList(String path, HashMap<String, List<String>> fileList)
	{
        File root = new File( path );
        File[] list = root.listFiles();
        for ( File f : list ) {
            if ( f.isDirectory() ) {
            	//if (f.getName().startsWith("NULL") || f.getName().matches("((-|\\+)?[0-9]+(\\.[0-9]+)?)+"))
            	//if (f.getName().matches("((-|\\+)?[0-9]+(\\.[0-9]+)?)+"))
            		GetFileList(f.getAbsolutePath(), fileList);
            }
            else {
            	String classificationFolder = f.getAbsolutePath().split("\\\\")[f.getAbsolutePath().split("\\\\").length - 2];
            	//if (classificationFolder.matches("((-|\\+)?[0-9]+(\\.[0-9]+)?)+"))
            	{
            		if (fileList.containsKey(classificationFolder) == false)
            			fileList.put(classificationFolder, new ArrayList<String>());
            		fileList.get(classificationFolder).add(f.getAbsolutePath());
            	}
            }
        }
        return fileList;
	}
}
